import gymnasium as gym
import highway_env
from stable_baselines3 import DQN
from stable_baselines3.common.evaluation import evaluate_policy


if __name__ == '__main__':
    # Create environment
    env = gym.make("highway-fast-v0", render_mode="rgb_array")

    # load model
    model = DQN.load("highway_dqn_model", env=env)

    mean_reward, std_reward = evaluate_policy(
        model,
        model.get_env(),
        deterministic=True,
        render=True,
        n_eval_episodes=10)

    print(mean_reward)
